package halberd.filter;

import halberd.DataType;
import halberd.ImgFilter;
import halberd.data.ImgData;
import halberd.data.ImgDataFactory;
import halberd.utils.RGBUtils;

/**
 * 白平衡
 * 
 * @author jiang
 *
 */
public class WhiteBalanceFilter implements ImgFilter {

	public ImgData process(ImgData input) {
		if (!canProcess(input)) {
			return null;
		}
		ImgData result = ImgDataFactory.createImgData(input.getWidth(), input.getHeight());
		whiteBalance(input.getDatas(), result.getDatas(), input.getWidth(), input.getHeight());
		result.setDataType(input.getDataType());
		return result;
	}

	private boolean canProcess(ImgData input) {
		if (input == null) {
			return false;
		}
		DataType dt = input.getDataType();
		if (dt == DataType.RGB) {
			return true;
		}
		if (dt == DataType.RGBA) {
			return true;
		}
		return false;
	}

	private void whiteBalance(int[][][] src, int[][][] dst, int w, int h) {

		int total = w * h;
		double avg_r = 0, avg_g = 0, avg_b = 0;
		double avg = 0;
		double gain_r, gain_g, gain_b;
		for (int i = 1; i < w; ++i) {
			for (int j = 1; j < h; ++j) {
				avg_r += src[i][j][0];
				avg_g += src[i][j][1];
				avg_b += src[i][j][2];
			}
		}
		avg_r = avg_r / total;
		avg_g = avg_g / total;
		avg_b = avg_b / total;
		avg = (avg_r + avg_g + avg_b) / 3;
		gain_r = avg / avg_r;
		gain_g = avg / avg_g;
		gain_b = avg / avg_b;
		for (int i = 1; i < w; ++i) {
			for (int j = 1; j < h; ++j) {
				int r = (int) (src[i][j][0] * gain_r);
				int g = (int) (src[i][j][1] * gain_g);
				int b = (int) (src[i][j][2] * gain_b);
				dst[i][j][0] = RGBUtils.clamp255(r);
				dst[i][j][1] = RGBUtils.clamp255(g);
				dst[i][j][2] = RGBUtils.clamp255(b);
				dst[i][j][3] = src[i][j][3];
			}
		}
	}
}
